from math import log, sqrt, exp, e
from collections import Iterable
from sys import stdout

from utils.misc import NAN, UNDEF
from utils.graph import tree
from utils.approx import Approx

from opt.treesearch.node import TreeNode
from opt import Infeasible


class MCTS_MetaNode(object):
    """
    An MCTS_MetaNode object attaches itself to a TreeNode object, encapsulating information that
    is relevant for Monte Carlo tree search.
    """
    __slots__ = ("solver", "node", "parent", "children", "next_branch",
                 "remaining_branches", "sim_count", "sim_fails", "sim_score",
                 "sim_best_score", "sim_total_score", "sim_total_infeas")

    def __init__(self, solver, node):
        self.solver = solver            # reference to an MCTS solver object
        self.node = node                # the TreeNode being wrapped by this meta node
        self.parent = None              # parent meta-node
        self.children = []              # list of active child meta-nodes
        self.next_branch = None         # next branch to expand
        self.remaining_branches = None  # branches awaiting expansion
        self.init_branches()

        self.sim_count = 0                # number of simulations run under this node
        self.sim_fails = 0                # number of failed simulations under this node
        self.sim_score = None             # objective value obtained in this node's simulation
        self.sim_best_score = Infeasible.INFINITY  # best (minimum) simulation score from this node
        self.sim_total_score = 0.0        # sum of simulation scores under this node
        self.sim_total_infeas = 0.0       # sum of infeasibilities from simulations under this node

    @property
    def is_expandable(self):
        """A node is expandable while it has expanded branches."""
        return self.remaining_branches is not None

    @property
    def is_exhausted(self):
        """A node is exhausted if it is completely expanded (remaining_branches is None) and all
        its children have been exhausted and thus removed from the tree (len(children) == 0)."""
        return self.remaining_branches is None and len(self.children) == 0

    def init_branches(self):
        if self.node.is_leaf():
            self.next_branch = None
            self.remaining_branches = None
        else:
            self.remaining_branches = iter(self.node.branches())
            self.advance_branch()

    def advance_branch(self):
        try:
            self.next_branch = self.remaining_branches.next()
        except StopIteration:
            self.next_branch = None
            self.remaining_branches = None
            self._finished_expansion()
            self.solver.channel.emit(SIGNALS.EXPANSION_COMPLETE)
            # remove this node from the tree if it is not the root node and it is exhausted
            if self.parent is not None and self.is_exhausted():
                self.parent._remove_child(self)

    # ----------------------------------------------------------------------------------------------
    # TREE MANAGEMENT FUNCTIONS
    def create_next_child(self):
        assert self.remaining_branches is not None
        # create child node and apply the next unexpanded branch to it
        child_node = self.node.copy()
        child_node.apply(self.next_branch)
        child = None
        if not self.solver._can_prune(child_node):
            child = type(self)(self.solver, child_node)
            # if the child is a leaf containing a feasible solution, we must *discard* its score
            # from the tree;  otherwise (child is not a leaf) we attach the node to the tree
            if child.remaining_branches is None:
                self._discard_sim_score(child_node.objective())
            else:
                child.parent = self
                self.children.append(child)
        # advance 'next_branch' to the next unexpanded branch
        try:
            self.next_branch = self.remaining_branches.next()
        except StopIteration:
            # finished expanding the node
            self.next_branch = None
            self.remaining_branches = None
            self._finished_expansion()
            # remove this node from the tree if it is not the root node and it is exhausted
            if self.parent is not None and self.is_exhausted():
                self.parent._remove_child(self)
        return child

    def _finished_expansion(self):
        """We have generated all of the node's children, so from now on we won't need the actual
        node data anymore; the metanode will suffice. This is useful for releasing memory."""
        self.node = None

    def _remove_child(self, child):
        parent = self
        while True:
            assert len(child.children) == 0
            # unlink the child from the tree
            parent.children.remove(child)
            child.parent = None
            # if the child has a simulation, we must subtract and discard its value
            if child.sim_count > 0:
                assert child.sim_count == 1
                parent._discard_sim_score(child.sim_score)
                parent._subtract_sim_score(child.sim_score)
            # if this node just became empty, we recursively remove it from its parent
            if len(parent.children) == 0 and parent.parent is not None:
                child = parent
                parent = parent.parent
            else:
                break

    def chop(self):
        """Remove a node and the whole subtree under it from the tree."""
        stack = [self]
        while len(stack) > 0:
            meta = stack.pop()
            if len(meta.children) == 0:
                meta.parent._remove_child(meta)
            else:
                stack.extend(meta.children)

    def do_simulation(self, sim_result=UNDEF):
        """Complete the solution in this metanode's node by calling the simulation() method on the
        node. The simulation() method is responsible for completing the partial solution in the
        node and either returning a leaf node (which should be a new node, created using copy()),
        or simply the objective value of the simulation.  If a leaf is returned by simulation(),
        it is returned to the solver to check it for improvement."""
        assert self.sim_count == 0
        if sim_result is UNDEF:
            sim_result = self.node.simulation()
        z = sim_result.objective() if isinstance(sim_result, TreeNode) else sim_result
        self._set_sim_score(z)
        if isinstance(z, Infeasible):
            solver = self.solver
            if z > solver.max_infeas and z.degree != INFINITY:
                solver.max_infeas = z
            if z < solver.min_infeas:
                solver.min_infeas = z
        return sim_result

    def _set_sim_score(self, score):
        """Increment the number of simulations and set the simulation score for this node.  The
        score is backpropagated, updating any parents of 'self' whose best score is worse than
        'score'."""
        assert score is not None
        self._add_sim_count(+1)
        self.sim_score = score
        if isinstance(score, Infeasible):
            self._add_sim_fails(+1)
            meta = self
            degree = score.degree
            while meta is not None:
                meta.sim_total_infeas += degree
                if meta.sim_best_score > score:
                    meta.sim_best_score = score
                meta = meta.parent
        else:
            meta = self
            while meta is not None:
                meta.sim_total_score += score
                if meta.sim_best_score > score:
                    meta.sim_best_score = score
                meta = meta.parent

    def _subtract_sim_score(self, score):
        """When a simulation is removed from the tree, its score must be subtracted from the total
        score of each meta-node in the path up to the root.  This method also decrements the node's
        simulation count."""
        assert score is not None
        self._add_sim_count(-1)
        if isinstance(score, Infeasible):
            self._add_sim_fails(-1)
            meta = self
            degree = score.degree
            while meta is not None:
                meta.sim_total_infeas -= degree
                meta = meta.parent
        else:
            meta = self
            while meta is not None:
                meta.sim_total_score -= score
                meta = meta.parent

    def _discard_sim_score(self, score):
        """Discard a simulation score from the tree.  Each of node's ancestors must discard its own
        simulation if its score is equal to 'score'.  Also, if 'score' is equal to an ancestor's
        best score, that ancestor must recompute its best score, after which three cases may occur:
            1) best score stays the same (if a sibling is able to reach the same score)
            2) best score becomes worse (if no sibling has best score equal to 'score')
            3) best score becomes INFEASIBLE (if there are no siblings with simulations)
        """
        assert score is not None
        meta = self
        while meta is not None:
            # if the node's own simulation score could possibly have been obtained through
            # the child associated with 'score', we discard this node's simulation score
            # and decrement its simulation count by one
            if score == meta.sim_score:
                meta.sim_score = None
                meta._subtract_sim_score(score)
            # recompute the node's best score from the node's own simulation score
            # and the best scores of the node's existing children
            if score == meta.sim_best_score:
                cand_scores = [c.sim_best_score for c in meta.children]
                if meta.sim_score is None:
                    cand_scores.append(Infeasible.INFINITY)
                else:
                    cand_scores.append(meta.sim_score)
                meta.sim_best_score = min(cand_scores)
            # proceed up the tree until we reach the root
            meta = meta.parent

    def _add_sim_count(self, n):
        """Adds 'n' to this node's 'sim_count' and backpropagates up to the root."""
        meta = self
        while meta is not None:
            meta.sim_count += n
            meta = meta.parent

    def _add_sim_fails(self, n):
        """Adds 'n' to this node's 'sim_fails' and backpropagates up to the root."""
        meta = self
        while meta is not None:
            meta.sim_fails += n
            meta = meta.parent

    # ----------------------------------------------------------------------------------------------
    # SELECTION-RELATED FUNCTIONS
    def exploration_scores_plain(self):
        """Compute a list with the exploration scores of this node's children."""
        if self.sim_count == 0:
            return [0.0] * len(self.children)
        exploration_coeff = self.solver._exploration_coeff
        two_log_parent_sims = 2.0 * log(self.sim_count)
        scores = []
        for child in self.children:
            if child.sim_count == 0:
                scores.append(INFINITY)
            else:
                scores.append(exploration_coeff * sqrt(two_log_parent_sims / child.sim_count))
        return scores

    def exploration_scores_weighted(self):
        """Like exploration_scores_plain(), but the exploration term is multiplied by the mean X_n
        (i.e. computed using the solver's _exploitation_score_fnc() on the mean simulation score
        instead of the best score)."""
        solver = self.solver
        z_star = float(solver.upper_bound)
        w_star = float(solver.worst_value)
        delta = (w_star - z_star) if w_star != z_star else 1.0
        exploitation_score_fnc = solver._exploitation_score_fnc
        scores = self.exploration_scores_plain()
        for i, child in enumerate(self.children):
            sim_successes = child.sim_count - child.sim_fails
            if sim_successes == 0:
                mean_reward = 1.0
            else:
                mean_sim_score = child.sim_total_score / sim_successes
                mean_reward = (w_star - mean_sim_score) / delta
                if mean_reward == 0.0:
                    mean_reward = 1.0  # prevent the worst branch from becoming excluded forever
            scores[i] *= exploitation_score_fnc(mean_reward)
        return scores

    def exploitation_scores(self):
        """Compute a list with the exploitation scores of this node's children."""
        solver = self.solver
        z_star = solver.upper_bound
        if z_star is INFINITY:
            return [0.0] * len(self.children)
        w_star = solver.worst_value
        delta = w_star - z_star
        if delta == 0.0:
            return [0.0 if isinstance(c.sim_best_score, Infeasible) else 1.0 for c in self.children]
        exploitation_score_fnc = self.solver._exploitation_score_fnc
        scores = []
        for child in self.children:
            z_child = child.sim_best_score
            if isinstance(z_child, Infeasible):
                scores.append(0.0)
            else:
                scores.append(exploitation_score_fnc((w_star - z_child) / delta))
        return scores

    @staticmethod
    def exploitation_score_exponential(reward):
        """Exploitation score functions take a raw linear reward in [0, 1], and should return the
        exploitation score associated with that reward. The exploitation score should ideally also
        be in [0, 1]. Use values outside this range with caution!"""
        return (exp(reward) - 1.0) / (e - 1.0)

    @staticmethod
    def exploitation_score_linear(reward):
        """See exploitation_score_exponential()."""
        return reward

    def infeasibility_penalties(self):
        """Score term for penalizing failed simulations."""
        max_infeas = self.solver.max_infeas.degree
        min_infeas = self.solver.min_infeas.degree
        delta_infeas = max_infeas - min_infeas
        if delta_infeas == 0.0:
            if self.sim_count == self.sim_fails:
                return [0.0] * len(self.children)
            return [-c.sim_fails / c.sim_count for c in self.children]
        penalties = []
        for child in self.children:
            z_child = child.sim_best_score
            if isinstance(z_child, Infeasible):
                feas_score = (max_infeas - z_child.degree) / delta_infeas
                penalties.append((feas_score - 1.0) * child.sim_fails / child.sim_count)
            else:
                penalties.append(0.0)
        return penalties

    def selection_candidates(self):
        score_terms = [term(self) for term in self.solver._selection_score_terms]
        children = self.children
        max_score = -INFINITY
        candidates = []
        for i in xrange(len(children)):
            score_i = sum(term[i] for term in score_terms)
            if score_i > max_score:
                max_score = score_i
                candidates = [children[i]]
            elif score_i == max_score:
                candidates.append(children[i])
        return candidates

    def selection_candidates_debug(self):
        print "selecting child at", self
        cand = []
        max_score = -INFINITY
        for child, terms in zip(self.children, self.selection_terms()):
            score = sum(terms)
            print terms, score, child,
            if score > max_score:
                print "new best candidate",
                max_score = score
                cand = [child]
            elif score == max_score:
                print "tied with best",
                cand.append(child)
        print "best candidates:", cand, "\n"
        return cand

    def selection_terms(self):
        """Used by report() to show the values of individual scoring terms."""
        return zip(*(term(self) for term in self.solver._selection_score_terms))

    def selection_scores(self):
        """Simply the sum of all selection score terms."""
        return map(sum, self.selection_terms())

    # ----------------------------------------------------------------------------------------------
    # AUXILIARY INSPECTION AND DEBUGGING FUNCTIONS AND PROPERTIES
    @property
    def sim_mean_score(self):
        """Average score of successful simulations run under this node."""
        sim_successes = self.sim_count - self.sim_fails
        return NAN if sim_successes == 0 else self.sim_total_score / sim_successes

    @property
    def sim_mean_infeas(self):
        """Average infeasibility of failed simulations run under this node."""
        return NAN if self.sim_fails == 0 else self.sim_total_infeas / self.sim_fails

    @property
    def path(self):
        """The sequence of nodes from the root of the tree to 'self'."""
        path = [self]
        parent = self.parent
        while parent is not None:
            path.append(parent)
            parent = parent.parent
        path.reverse()
        return path

    @property
    def dimensions(self):
        """Return the size (# of nodes) and height (# of levels) of the tree rooted at 'self'."""
        if len(self.children) == 0:
            return (1, 1)
        size = 0
        height = 0
        for child in self.children:
            child_size, child_height = child.dimensions
            size += child_size
            height = max(height, child_height)
        size += 1
        height += 1
        return (size, height)

    @property
    def depth(self):
        """# of nodes in the path between this node (exclusive) and the root (inclusive)."""
        depth = 0
        parent = self.parent
        while parent is not None:
            depth += 1
            parent = parent.parent
        return depth

    @property
    def mean_depth(self):
        """Mean depth of *all* nodes of the tree."""
        total_depth = 0
        total_nodes = 0
        stack = [(self, 0)]
        while len(stack) > 0:
            meta, depth = stack.pop()
            total_depth += depth
            total_nodes += 1
            depth += 1
            stack.extend((child, depth) for child in meta.children)
        return float(total_depth) / total_nodes

    @property
    def mean_leaf_depth(self):
        """Mean depth of the *leaf* nodes of the tree. Note that leaf here does not necessarily
        refer to a node containing an actual complete solution, but a node which has no children
        at the time of the call."""
        total_depth = 0
        total_leaves = 0
        stack = [(self, 0)]
        while len(stack) > 0:
            meta, depth = stack.pop()
            if len(meta.children) > 0:
                depth += 1
                stack.extend((child, depth) for child in meta.children)
            else:
                total_depth += depth
                total_leaves += 1
        return float(total_depth) / total_leaves

    def __len__(self):
        """Number of children of this node."""
        return len(self.children)

    def __getitem__(self, i):
        """meta[i] is short for meta.children[i]"""
        if isinstance(i, int):
            return self.children[i]
        elif isinstance(i, Iterable):
            meta = self
            for j in i:
                meta = meta.children[j]
            return meta
        else:
            raise TypeError("expecting integer or iterable")

    def to_networkx(self):
        get_children = lambda x: x.children
        get_sim_best_score = lambda x: (x.sim_best_score
                                        if not isinstance(x.sim_best_score, Infeasible)
                                        else None)
        return tree.build(root=self, get_children=get_children, get_weight=get_sim_best_score)

    def draw(self, highlight_paths=(), axes=None, show=True):
        G = self.to_networkx()
        get_children = lambda x: x.children
        return tree.draw(G, root=self, get_children=get_children,
                         highlight_paths=highlight_paths,
                         axes=axes, show=show)

    def report(self, levels=1, indent="", ostream=stdout):
        self._report_line(ostream, indent, "\n")
        self._report_descendants(levels, "    "+indent, ostream)

    def _report_descendants(self, levels=1, indent="", ostream=stdout):
        if levels >= 1:
            for child, terms in zip(self.children, self.selection_terms()):
                score = sum(terms)
                terms = "".join("%+.03f" % term for term in terms)
                child._report_line(ostream, indent, " score=%.03f (%s)\n" % (score, terms))
                child._report_descendants(levels-1, "    "+indent, ostream)

    def _report_line(self, ostream, prefix="", suffix="\n"):
        line_fmt = self.report_line_fmt
        z = format_sim_score(self.sim_score)
        z_mean = format_sim_score(self.sim_mean_score)
        z_star = format_sim_score(self.sim_best_score)
        ostream.write(line_fmt % (prefix, self.depth, z, z_mean, z_star,
                                  self.sim_count, self.sim_fails, suffix))

    report_line_fmt = "%sdepth=%d z=%s z'=%s z*=%s sims=%d (%d fails)%s"

    def iter_depth_first(self):
        """Iterate through all the nodes in the tree in depth-first order.
        Returns '(metanode, depth)' tuples."""
        stack = [(self, 0)]
        while len(stack) > 0:
            meta, depth = stack.pop()
            yield meta, depth
            depth += 1
            stack.extend((child, depth) for child in meta.children)

    def validate(self):
        for meta, _ in self.iter_depth_first():
            # verify that the sim_best_score of meta is indeed the best score among
            # its own simulation and its children's best simulations
            scores = [c.sim_best_score for c in meta.children]
            if meta.sim_score is None:
                scores.append(Infeasible.INFINITY)
            else:
                scores.append(meta.sim_score)
            assert meta.sim_best_score == min(scores)
            # the total simulation count at meta should be equal to the sum of the simulations
            # done by its children plus 1 if meta's own simulation hasn't been discarded yet
            own_sim = 0
            own_sim_score = 0.0
            if meta.sim_score is not None:
                own_sim = 1
                if not isinstance(meta.sim_score, Infeasible):
                    own_sim_score = meta.sim_score
            assert meta.sim_count - own_sim == sum(c.sim_count for c in meta.children)
            assert (Approx(meta.sim_total_score) - own_sim_score ==
                    sum(c.sim_total_score for c in meta.children))
            # verify if all children are correctly linked to the parent
            assert all(child.parent is meta for child in meta.children)


def format_sim_score(score):
    return ("undef" if score is None else
            "infeas({:.03f})".format(score.degree) if isinstance(score, Infeasible) else
            "{:.03f}".format(score))
